package ge

import (
	"bytes"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/sirupsen/logrus"
)

// bodyLogWriter 定义一个存储响应内容的结构体
//
// 在bodyLogWriter结构体中封装了gin的responseWriter, 然后在重写的Write方法中, 首先向bytes.Buffer中写数据,然后响应
//
// 这保证了我们可以正确的获取到响应内容最后就是中间件的实现, 其中最重要的一点就是用我们自定义的bodyLogWriter来代替ctx.Writer, 保证响应会保留一份在bytes.Buffer中
type bodyLogWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

// Write 读取响应数据
//
//	@receiver w
//	@param b
//	@return int
//	@return error
func (w bodyLogWriter) Write(b []byte) (int, error) {
	w.body.Write(b)
	return w.ResponseWriter.Write(b)
}

// ResponseLog gin响应的日志中间件
//
//	@param ctx
func ResponseLog(ctx *gin.Context) {
	// 初始化bodyLogWriter
	blw := &bodyLogWriter{
		body:           bytes.NewBufferString(""),
		ResponseWriter: ctx.Writer,
	}
	ctx.Writer = blw
	ctx.Next()
	tp := blw.ResponseWriter.Header().Get("Content-Type")
	if strings.Contains(tp, "json") || strings.Contains(tp, "text") {
		// 响应内容
		responseBody := blw.body.String()
		if len(responseBody) > 150 {
			responseBody = responseBody[:150] // 限制打印数量
		}
		logrus.Info(responseBody)
	}
}
